import numpy as np
import platgo as pg


class DTLZ8(pg.Problem):

    def __init__(self, M: int = 3) -> None:
        self.name = 'DTLZ8'
        self.type['multi'], self.type['many'], self.type['real'], self.type['large'], self.type['expensive'], self.type['constrained'] = [
                                                                                                                   True] * 6
        self.M = M
        self.D = M * 10
        lb = [0] * self.D
        ub = [1] * self.D
        self.borders = np.array([lb, ub])
        super().__init__()

    def cal_obj(self, pop: pg.Population) -> None:
        decs = pop.decs
        pop.objv = np.zeros((pop.N, self.M))
        pop.cv = np.zeros((pop.N, self.M))
        for i in range(1, self.M + 1):
            f = np.sum(decs[:, ((i - 1) * 10):i * 10], axis=1, keepdims=True) * 0.1
            pop.objv[:, i - 1:i] = f[:, :]
        pop.cv[:, :(self.M - 1)] = (pop.objv[:, -1:] + 4 * pop.objv[:, :(self.M - 1)] - 1) * -1
        if self.M == 2:
            pop.cv[:, -1:] = 0
        else:
            minValue = np.sort(pop.objv[:, :(self.M - 1)], axis=1)
            pop.cv[:, -1:] = (2 * pop.objv[:, -1:] + np.sum(minValue, axis=1, keepdims=True) - 1) * -1

    def get_optimal(self) -> np.ndarray:
        # 目标空间均匀采样函数还未完成
        raise NotImplementedError("get optimal has not been implemented")


if __name__ == '__main__':
    d = DTLZ8(3)
    pop = d.init_pop(5)
    # pop = pg.Population(decs=np.array([[2, 2.0, 2], [1, 2, 2], [1, 3, 2]]))
    # pop = pg.Population(decs=np.random.uniform(0, 1, (10, 30)))
    print(d.borders)
    print(pop.decs)
    d.cal_obj(pop)  # 计算目标函数值
    print(pop.objv)
    print(pop.cv)
    print(d.type["constrained"])

